/**
 * 
 */
package com.gragra.sampling.nonproj.structures.projective.singleroot;
import java.util.Arrays;
import java.util.BitSet;
import com.gragra.sampling.PopulationProbabilityAssigner;
import com.gragra.sampling.ThreadedUniformAccess;
import com.gragra.sampling.binary.ConvertSample;
import com.gragra.sampling.nonproj.ProposalDependencyAssigner;
import com.gragra.sampling.sequentialStructure.WordAndTagHolder;
/**
 * @author Christoph Teichmann
 * 11.10.2013
 */
public class GreedySingleRootProjective extends GibbsMHSingleRootProjective
{
	/**
	 * 
	 * @param forCandidates
	 * @param thread
	 * @param tua
	 * @return
	 */
	protected int[] produceShuffle(BitSet forCandidates, int thread, ThreadedUniformAccess tua)
	{
		int[] container = shuffler[thread];
		Arrays.fill(container, -1);
		int pos = 0;
		if(tua.nextBoolean(thread))
		{
			for(int head=forCandidates.nextSetBit(0);head>=0;head=forCandidates.nextSetBit(head+1))
			{container[pos++] = head;}
		}
		else
		{
			for(int head=forCandidates.previousSetBit(forCandidates.size()+1);head>=0;head=forCandidates.previousSetBit(head-1))
			{container[pos++] = head;}
		}
		return container;
	}
	/**
	 * 
	 */
	private final int[][] shuffler;
	/**
	 * 
	 * @param thread
	 * @return
	 */
	protected void restore(int thread)
	{
		for(int pos=0;pos<this.getLength();++pos)
		{attachSingleFromStorage(pos, thread);}
	}
	/**
	 * @param thread
	 * @param pos
	 */
	protected void attachSingleFromStorage(int pos, int thread)
	{
		if(this.getHeadPosition(pos, thread) == -1)
		{this.attach(pos, this.getHeadFromStorage(pos, thread), thread);}
		else
		{this.setNewHead(pos, this.getHeadFromStorage(pos, thread), thread);}
	}
	/**
	 * 
	 * @param forPos
	 * @param thread
	 * @return
	 */
	protected int getHeadFromStorage(int forPos, int thread)
	{return this.storage[thread][forPos];}
	/**
	 * 
	 * @param thread
	 */
	protected void writeToStorage(int thread)
	{
		int[] sto = storage[thread];
		for(int i=0;i<this.getLength();++i)
		{sto[i] = this.getHeadPosition(i, thread);}
	}
	/**
	 * 
	 */
	private final int[][] storage;
	/**
	 * 
	 * @param position
	 * @param thread
	 * @return
	 */
	protected BitSet getWorkingSet(int position, int thread)
	{
		BitSet ret = considered[thread];
		ret.clear();
		this.fillContainerWithChildren(position, ret, thread);
		return ret;
	}
	/**
	 * 
	 */
	private final BitSet[] considered;
	/**
	 * @param wth
	 * @param threads
	 * @param cs
	 */
	public GreedySingleRootProjective(WordAndTagHolder wth, int threads,
			ConvertSample cs)
	{
		super(wth, threads, cs);
		considered = new BitSet[threads];
		for(int i=0;i<considered.length;++i)
		{considered[i] = new BitSet();}
		this.storage = new int[threads][this.getLength()];
		this.shuffler = new int[threads][this.getLength()+1];
	}
	/**
	 * @param pos
	 * @param pa
	 * @param annealing
	 * @param thread
	 * @param tua
	 * @throws IllegalStateException
	 */
	protected void greedySamplingForPosition(int pos,
			PopulationProbabilityAssigner pa, double annealing, int thread,
			ThreadedUniformAccess tua) throws IllegalStateException
	{
		double oldProb = this.makeLogProbForThisConsideringAnnealing(pa, annealing, thread);
		BitSet relevant = this.getWorkingSet(pos, thread);
		if(relevant.cardinality() < 2)
		{return;}
		double correction = 0.0;
		ProposalDependencyAssigner pda = (ProposalDependencyAssigner) pa;
		this.writeToStorage(thread);
		this.clear(relevant,thread);
		int[] shuffling = this.produceShuffle(relevant, thread, tua);
		for(int i=0;shuffling[i] != -1 && i<shuffling.length;++i)
		{
			int toAssign = shuffling[i];
			BitSet candidates = makeCandidates(toAssign, pos, thread, relevant, pda);
			double[] probs = makeProbs(toAssign, candidates, annealing, pda, thread);
			correction += Math.log(probs[this.getHeadFromStorage(toAssign, thread)]);
			this.attachSingleFromStorage(toAssign, thread);
		}
		this.clear(relevant, thread);
		for(int i=0;shuffling[i] != -1 && i<shuffling.length;++i)
		{
			int toAssign = shuffling[i];
			BitSet candidates = makeCandidates(toAssign, pos, thread, relevant, pda);
			double[] probs = makeProbs(toAssign, candidates, annealing, pda, thread);
			int choice = -1;
			double border = tua.nextDouble(thread);
			for(int head=candidates.nextSetBit(0);head>=0;head=candidates.nextSetBit(head+1))
			{
				border -= probs[head];
				if(border <= 0.0)
				{
					choice = head;
					break;
				}
			}
			if(choice < 0)
			{throw new IllegalStateException("Probabilities did not sum to one.");}
			this.attach(toAssign, choice, thread);
			correction -= Math.log(probs[choice]);
		}
		double newProb = this.makeLogProbForThisConsideringAnnealing(pa, annealing, thread);
		if(Math.log(tua.nextDouble(thread)) > newProb-oldProb+correction)
		{this.restore(thread);}
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.nonproj.structures.nonprojective.NonProjDepdTree#resample(com.gragra.sampling.PopulationProbabilityAssigner, double, int, com.gragra.sampling.ThreadedUniformAccess)
	 */
	@Override
	public void resample(PopulationProbabilityAssigner pa, double annealing,
			int thread, ThreadedUniformAccess tua)
	{
		this.greedySamplingForPosition(this.getRootPosition(), pa, annealing, thread, tua);
		super.resample(pa, annealing, thread, tua);
	}
	/**
	 * 
	 * @param toAssign
	 * @param candidates
	 * @param annealing
	 * @param pda
	 * @param thread
	 * @return
	 */
	private double[] makeProbs(int toAssign, BitSet candidates,
			double annealing, ProposalDependencyAssigner pda, int thread)
	{
		double max = Double.NEGATIVE_INFINITY;
		double[] ps = this.getProbsHolder(thread);
		for(int head=candidates.nextSetBit(0);head>=0;head=candidates.nextSetBit(head+1))
		{
			this.attach(toAssign, head, thread);
			max = Math.max(max, ps[head] = makeProposalWithAnnealing(toAssign, annealing, pda, thread, head));
			this.detach(toAssign, thread);
		}
		double sum = 0.0;
		for(int head=candidates.nextSetBit(0);head>=0;head=candidates.nextSetBit(head+1))
		{sum += ps[head] = Math.exp(ps[head]-max);}
		for(int head=candidates.nextSetBit(0);head>=0;head=candidates.nextSetBit(head+1))
		{ps[head] = ps[head]/sum;}
		return ps;
	}
	/**
	 * @param toAssign
	 * @param annealing
	 * @param pda
	 * @param thread
	 * @param head
	 * @return
	 */
	protected double makeProposalWithAnnealing(int toAssign, double annealing,
			ProposalDependencyAssigner pda, int thread, int head)
	{return pda.makeProposalLogProb(toAssign, head, this, thread)*annealing;}
	/**
	 * @param pos
	 * @param thread
	 * @param relevant
	 * @param pda
	 * @return
	 */
	protected BitSet makeCandidates(int pos,int cover, int thread, BitSet relevant,
			ProposalDependencyAssigner pda)
	{
		BitSet candidates = this.makeAllHeadCandidates(pos, thread, pda);
		boolean in = candidates.get(cover);
		candidates.and(relevant);
		if(in)
		{candidates.set(cover);}
		return candidates;
	}
	/**
	 * 
	 * @param relevant
	 */
	protected void clear(BitSet relevant, int thread)
	{
		for(int head=relevant.nextSetBit(0);head>=0;head=relevant.nextSetBit(head+1))
		{this.detach(head, thread);}
	}
}